#define CPU_ONLY
#define USE_OPENCV
#include <google/protobuf/text_format.h>

#include <lmdb.h>

#include <cstdint>
#include <sys/stat.h>

#include <fstream>
#include <string>

#include "caffe/proto/caffe.pb.h"
#include "caffe/util/db.hpp"
#include "caffe/util/format.hpp"
#include "caffe/util/io.hpp"
#include <amber/helper.hpp>
#include <opencv2/opencv.hpp>
#include <vector>

using namespace caffe;
using std::string;
using std::vector;

void convert_dataset(const char* imageListFileName,
    const char* db_path, const string& db_backend)
{
    // 打开LMDB数据库
    std::unique_ptr<db::DB> db(db::GetDB(db_backend));
    db->Open(db_path, db::NEW);
    // 启动事务
    std::unique_ptr<db::Transaction> txn(db->NewTransaction());

    // NCHW
    // 创建Datum对象，表示一条训练数据
    Datum datum;
    // 设置训练数据的通道/高度/宽度

    int32_t count = 0;
    std::string value;

    // 打开图片列表
    // 图片列表文件中每一行是一个图片的路径
    vector<string> imageListFile = amber::listFiles(imageListFileName);

    for (auto file : imageListFile) {

        cv::Mat matrix = cv::imread(file, CV_LOAD_IMAGE_COLOR);
        CVMatToDatum(matrix, &datum);

        datum.set_label(count);

        // 生成数据库中的key（每一条数据的唯一标识符）
        datum.SerializeToString(&value);
        // 在事务中写入数据
        txn->Put(caffe::format_int(count, 8), value);

        if (count % 1000 == 0) {
            std::cout << "Commit " << count << " files." << std::endl;
            // 每1000条数据提交一次事务
            txn->Commit();
        }

        count++;
    }

    db->Close();
}

int main(int argc, char** argv)
{
    if (argc < 3) {
        std::cerr << "Invalid arguments" << std::endl;

        std::cout << "This script converts the images to\n"
                     "the lmdb format used by Caffe to load data.\n"
                     "Usage:\n"
                     "    convert-image-data input_image_file_list output_db_file\n"
                  << std::endl;
        return EXIT_FAILURE;
    }

    // 第一个参数是图片列表文件路径
    // 第二个参数是输出数据库路径
    convert_dataset(argv[1], argv[2], "lmdb");

    return 0;
}
